Add EntropyGatedChunkKVPress - #263
Conversation
Co-authored-by: Liran Azran <liran.azr90@gmail.com> Signed-off-by: Shahar Ben-Ishay <shahar.benishay@gmail.com>
|
Some further results. All runs used the same model and same hyperparams ( RULER-4096, mean over the 13 subtasks:
The gate is not tied to one scorer. Swapping the inner
Thanks! |
|
/ok to test 7797071 |
SimJeg
left a comment
There was a problem hiding this comment.
Hi, thanks for the PR! You will find a first batch of comments to improve the readability of this new press
|
Hi @ShaharBenIshay @Liranitz, thanks for the first updates. Tell me when you think you're done with this first batch of comments. |
…istry cleanup - kvpress/presses/entropy_gated_chunkkv_press.py: subclass ChunkKVPress instead of BasePress, dropping the inherited press field, __post_init__, post_init_from_model, and compression_ratio property/setter; keep chunk_length default at 10. Remove the redundant explanatory comments and hoist the epsilon to a module-level EPSILON constant. - tests/presses/test_entropy_gated_chunkkv_press.py: deleted. The dedicated test duplicated coverage already provided by the test_presses_run wrapper matrix. - tests/presses/test_presses.py: removed the redundant test_entropy_gated_chunkkv_press function for the same reason; the press stays covered via the EntropyGatedChunkKVPress entry in the wrapper_press matrix. - evaluation/evaluate_registry.py: dropped the explicit chunk_length=10, rescue_size=4 from the registry entry; both are defaults on the press now, so EntropyGatedChunkKVPress(press=SnapKVPress()) is enough. - README.md: tightened the one-line description to match the other press entries. Signed-off-by: Liran Azran <liran.azr90@gmail.com>
Signed-off-by: Liran Azran <liran.azr90@gmail.com>
…eanup Negative-score correctness: - Drop clamp(min=0) on the per-token scores; ranking, median, argsort and all top-k selection now use the raw scores, so signed scorers (e.g. KeyDiffPress) are ordered correctly instead of silently collapsing to ties. - Compute the within-chunk entropy from a per-chunk min-shift (subtract the chunk minimum only when it is negative) so the scores form a valid distribution. Remove the entropy_threshold kwarg: - Always use the per-example median entropy as the spikiness cutoff, drop the entropy_threshold field, its docstring entry. Guard chunk_length and simplify: - Assert chunk_length > 1 in __post_init__ and remove the c == 1 entropy branch it makes unreachable (the length-1 partial-chunk case is still handled separately). - Remove the redundant budget >= kv_len early return. Naming and readability: - Rename locals to intent-revealing names (chunk_len, scores, chunk_token_scores, chunk_scores, chunk_entropy, score_threshold, high_score_chunks, low_entropy_chunks, low_entropy_chunk_length) and restructure the greedy loop around a single n_kept quantity. Signed-off-by: Liran Azran <liran.azr90@gmail.com>
Signed-off-by: Liran Azran <liran.azr90@gmail.com>
9f34142 to
21f6ed6
Compare
|
Thank you @SimJeg for the meaningful comments (and for your time), we have resolved all of them. If, after reviewing the revised version, there are any other comments, we would happily fix and resolve them too. |
SimJeg
left a comment
There was a problem hiding this comment.
Hi, thanks for the updates ! I left a few final comments
- Collapse the negative-score rebasing to a single out-of-place line, dropping the (chunk_token_scores < 0).any() guard and the chunk_min temporary. Keep it out-of-place: chunk_token_scores is a view into scores, so an in-place -= would mutate scores and corrupt the later top-k ranking. - Consolidate the two __post_init__ asserts into one enforcing chunk_length > low_entropy_chunk_length >= 1, adding the missing lower bound so an important-but-spiky chunk always keeps at least one token. - Add a blank line before the greedy-pass section and drop the unused S / H_tilde notation from the section-1 comment. Signed-off-by: Liran Azran <liran.azr90@gmail.com>
|
/ok to test c57ddd4 |
|
LGMT ! I'm running the CI/CD and if it passes I will merge. Feel free to open a small PR once you get a preprint out |
|
@SimJeg, much appreciated the feedback. |
PR description
Adds
EntropyGatedChunkKVPress(EG-ChunkKV), as proposed and approved in #261.ChunkKVPressscores each chunk by its aggregate token importance and then keeps or drops thatchunk as a whole. Because the score is a magnitude statistic, it cannot distinguish a coherent
chunk, whose importance is spread across its tokens, from a spiky one, where a single needle
token holds nearly all the mass and the remaining
chunk_length - 1tokens are filler. Both canscore identically, and for the spiky chunk, keeping it whole spends
chunk_lengthcache slots topreserve one useful token — which, under a fixed budget, evicts a chunk that would otherwise be kept.
EG-ChunkKV adds a second, shape statistic computed from the same per-token scores ChunkKV already
has: the normalized within-chunk Shannon entropy
H̃_i ∈ [0, 1], where 1 means coherent and 0 meansconcentrated in a few tokens. Important-but-spiky chunks are reduced to their top-
rescue_sizetokens rather than kept whole, and the freed budget is spent on further chunks. The number of
retained tokens is exactly
max(1, floor((1 - compression_ratio) * kv_len)), identical toChunkKVPress, so the two are budget-matched by construction.Selection procedure
S_i: the mean of the head-summed, non-negative token scores in chunki. This is the samestatistic
ChunkKVPressranks chunks by, so both methods order chunks identically.H̃_i: the normalized Shannon entropy of those same token scores within the chunk.Lthat ChunkKV uses and walk chunks greedily in decreasing
S_i, spending budget as we go:If the rescues leave budget unspent, a final fill adds the highest-scoring not-yet-retained tokens,
so exactly
Lpositions are kept and the retained indices are sorted back into original order.Both gating masks are computed vectorized and moved to CPU once, so the sequential budget walk reads
no GPU scalars per iteration — without that, the per-chunk synchronization made this measurably
slower than plain ChunkKV.
Results
Llama-3.1-8B-Instruct, SnapKV as the inner scorer(for a fair comparison vs chunkkv, although we are even stronger with expected attention),
chunk_length=10,rescue_size=4.cris thefraction of the KV cache removed. LongBench is the 16-task average; LOOGLE is ROUGE-L × 100. Higher
is better. ChunkKV is run at the same
chunk_length=10so the comparison isolates the entropy gate.The gain holds at every ratio and grows as the budget tightens, peaking at +1.35 at
cr=0.90. LOOGLE shows the same direction under aggressive compression (+0.42 and +0.46) but isslightly behind at
cr=0.70(−0.11), which is consistent with the mechanism: the wasted budget onlybecomes costly once slots are genuinely scarce.
Notes for review
docstring cites only the ChunkKV paper.
default_pressescheckbox is intentionally unticked. That list is instantiated ascls(compression_ratio=...), but this press requirespress=and exposescompression_ratioas adelegating property, so it cannot be constructed that way.
ChunkKVPressis absent from the listfor the same reason. Instead the press is added to the
wrapper_pressmatrix intests/presses/test_presses.py, following theChunkKVPressprecedent, which exercises it againstevery press in
default_presses.chunk_lengthdefaults to 10, not 20 as inChunkKVPress. Finer granularity gives the gatemore chunks to reallocate budget between — the table above shows the mechanism is markedly stronger
at 10 — and it is the configuration all reported numbers use (however we did measure our method at c=20 and it will be presented in our preprint in depth).
Checklist
Before submitting a PR, please make sure:
Tests are working (
make test)Code is formatted correctly (
make style, on errors try fix withmake format)Copyright header is included
All commits are signed-off using
git commit -s(new press)
mypress_press.pyis in thepressesdirectory(new press)
MyPressis in__init__.py(new press)
README.mdis updated with a 1 liner about the new press in the Available presses section(new press) New press is in the
default_presseslist intests/default_presses.py(new press) A docstring is provided that follows the same structure as the existing ones